import json
import pickle

import cv2
import numpy as np
from auxmodels.core.data import NNDataset

from .augmenter import Augmenter
from .encoder import CenterEncoder


class CenterDataset(NNDataset):
    def __init__(self, sample_file, augmenter: Augmenter, num_cls: int, input_size) -> None:
        super().__init__()
        with open(sample_file, mode="rb") as f:
            self.sample_list = pickle.load(f)
        # 数据增强操作
        self.augmenter = augmenter
        # 目标编码操作
        self.encoder = CenterEncoder(input_size, num_cls)

    def __len__(self):
        return len(self.sample_list)

    def _restore_data(self, index):
        sample = self.sample_list[index]
        image = cv2.imread(sample["filename"])
        label = np.array(sample["targets"], dtype=np.float32)
        return image, label

    def _augment(self, image, label):
        return self.augmenter.run(image, label)

    def _build_target(self, label):
        return self.encoder.run(label)
